{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c348641e",
   "metadata": {},
   "source": [
    "This is part 1 of a tutorial series. We recommend following them in order, starting with [Part 0: Welcome to `musica`](0.%20Welcome%20to%20MUSICA.ipynb)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1000001",
   "metadata": {},
   "source": [
    "# Chapman Cycle Box Model\n",
    "\n",
    "This tutorial implements a box model of the [Chapman cycle](https://en.wikipedia.org/wiki/Chapman_cycle) — a set of reactions describing stratospheric ozone chemistry.\n",
    "\n",
    "In previous tutorials, we used `micm` and `tuv-x` independently. Now, we're going to use them together to calculate photolysis rate constants and solve the Chapman mechanism for a real geographic location (Boulder, CO) over 12 hours beginning at 6:30 AM local time. The chemistry will be solved for 10 vertically stacked grid cells in the stratosphere (20–30 km), where ozone chemistry is most active."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1000002",
   "metadata": {},
   "source": [
    "## 1. Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1000003",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from datetime import datetime, time, timedelta\n",
    "from zoneinfo import ZoneInfo\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pvlib\n",
    "import ussa1976\n",
    "import xarray as xr\n",
    "\n",
    "import musica\n",
    "from musica.constants import GAS_CONSTANT\n",
    "from musica.mechanism_configuration import Parser\n",
    "from musica.micm.solver_result import SolverState\n",
    "from musica.tuvx import vTS1\n",
    "from musica.utils import find_config_path\n",
    "\n",
    "SECONDS_PER_HOUR = 3600"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1000004",
   "metadata": {},
   "source": [
    "## 2. Location and Simulation Time\n",
    "\n",
    "We simulate over Boulder, CO. The simulation starts at 6:30 Mountain Time so we capture the morning rise in photolysis rates as the sun climbs.\n",
    "\n",
    "All times are tracked as timezone-aware `datetime` objects. We use UTC internally and convert to local time for display."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1000005",
   "metadata": {},
   "outputs": [],
   "source": [
    "boulder = (40.01879858223568, -105.27492413846649)  # (latitude, longitude)\n",
    "boulder_tz = ZoneInfo(\"America/Denver\")\n",
    "\n",
    "today_local = datetime.now(boulder_tz).date()\n",
    "sim_local  = datetime.combine(today_local, time(6, 30), tzinfo=boulder_tz)\n",
    "sim_time    = sim_local.astimezone(ZoneInfo(\"UTC\"))\n",
    "\n",
    "print(f\"Simulation start (UTC):   {sim_time}\")\n",
    "print(f\"Simulation start (local): {sim_time.astimezone(boulder_tz)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1000006",
   "metadata": {},
   "source": [
    "## 3. Load `tuv-x` and Define a Helper to Compute Photolysis Rates\n",
    "\n",
    "We load the vTS1 `tuv-x` calculator once and keep it alive for the duration of the simulation. The photolysis reactions in the Chapman mechanism are available from this `tuv-x` configuration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1000007",
   "metadata": {},
   "outputs": [],
   "source": [
    "tuvx = vTS1.get_tuvx_calculator()\n",
    "\n",
    "tuv_path = find_config_path(\"tuvx\", \"ts1_tsmlt.json\")\n",
    "with open(tuv_path, 'r') as f:\n",
    "    data = json.load(f)\n",
    "alias_mappings = data.get('__CAM options', {}).get('aliasing', {}).get('pairs', {})\n",
    "\n",
    "print(f\"Loaded {len(alias_mappings)} photolysis rate alias mappings\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d6103b7",
   "metadata": {},
   "source": [
    "Next, let's create a convenience function we can call at each time step to get the photolysis rate constants, by:\n",
    "1. Using `pvlib` to compute the solar zenith angle (SZA) at Boulder for that moment.\n",
    "2. Running TUV-x with that SZA to get photolysis rate constants at every altitude.\n",
    "3. Mapping the TUV-x reaction labels to the `PHOTO.*` parameter names expected by the MICM Chapman mechanism using the alias table in the config file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1000008",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_tuv_rates(utc_time, num_grid_cells, start=1):\n",
    "    \"\"\"Return photolysis rate constants for `num_grid_cells` altitude cells starting at index `start`.\"\"\"\n",
    "    lat, lon = boulder\n",
    "    solpos = pvlib.solarposition.get_solarposition(time=utc_time, latitude=lat, longitude=lon)\n",
    "    sza_deg = solpos['zenith'].item()\n",
    "\n",
    "    tuv_rates = tuvx.run(sza=np.deg2rad(sza_deg), earth_sun_distance=1.0)\n",
    "    end = start + num_grid_cells\n",
    "\n",
    "    photolysis_rate_constants = {}\n",
    "    for mapping in alias_mappings:\n",
    "        label    = mapping['to']\n",
    "        scale    = mapping.get(\"scale by\", 1)\n",
    "        tuv_label = mapping['from']\n",
    "        rate = tuv_rates.sel(reaction=tuv_label).photolysis_rate_constants.values * scale\n",
    "        photolysis_rate_constants[f'PHOTO.{label}'] = rate[start:end]\n",
    "\n",
    "    return photolysis_rate_constants, tuv_rates"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1000009",
   "metadata": {},
   "source": [
    "## 4. Altitude Grid and Atmospheric Conditions\n",
    "\n",
    "The Chapman cycle is most relevant in the stratosphere (~20–50 km). We use TUV-x altitude grid cells 20–30 (i.e., the 20th through 29th edges), which correspond to approximately 20–30 km. Temperature and pressure at those altitudes come from the US Standard Atmosphere 1976 model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1000010",
   "metadata": {},
   "outputs": [],
   "source": [
    "start          = 20   # TUV-x altitude grid index (skip surface + lower troposphere)\n",
    "num_grid_cells = 10   # number of stratospheric grid cells\n",
    "\n",
    "# Compute initial photolysis rates and grab the altitude edges for our grid cells\n",
    "photolysis_rate_constants, tuv_rates = get_tuv_rates(sim_time, num_grid_cells, start)\n",
    "vertical_edges = tuv_rates.vertical_edge[start:start + num_grid_cells].data\n",
    "\n",
    "print(f\"Altitude range: {vertical_edges.min():.1f} – {vertical_edges.max():.1f} km\")\n",
    "\n",
    "# US Standard Atmosphere T and P at those altitudes (km -> m)\n",
    "atm         = ussa1976.compute(z=vertical_edges * 1000, variables=[\"t\", \"p\"])\n",
    "temperature = atm['t'].values\n",
    "pressure    = atm['p'].values\n",
    "\n",
    "print(f\"Temperature range: {temperature.min():.1f} – {temperature.max():.1f} K\")\n",
    "print(f\"Pressure range:    {pressure.min():.1f} – {pressure.max():.1f} Pa\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1000011",
   "metadata": {},
   "source": [
    "## 5. Load the Chapman Mechanism and Create the Solver"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1000012",
   "metadata": {},
   "outputs": [],
   "source": [
    "parser    = Parser()\n",
    "mechanism = parser.parse(find_config_path(\"v1\", \"chapman\", \"config.json\"))\n",
    "\n",
    "solver = musica.MICM(\n",
    "    mechanism=mechanism,\n",
    "    solver_type=musica.SolverType.rosenbrock_standard_order\n",
    ")\n",
    "state = solver.create_state(num_grid_cells)\n",
    "\n",
    "print(f\"Chapman mechanism species: {list(state.get_species_ordering().keys())}\")\n",
    "print(f\"User-defined rate parameters: {list(state.get_user_defined_rate_parameters().keys())}\")\n",
    "for reaction in mechanism.reactions:\n",
    "    print(f\"[{reaction.type.name}] {reaction.to_equation()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1000013",
   "metadata": {},
   "source": [
    "## 6. Set Initial Concentrations\n",
    "\n",
    "Species concentrations are specified as volume mixing ratios and converted to mol m⁻³ using the ideal gas law:\n",
    "\n",
    "$$[X] = \\frac{\\chi_X \\cdot P}{R T}$$\n",
    "\n",
    "where $\\chi_X$ is the mixing ratio, $P$ pressure (Pa), $R$ the universal gas constant, and $T$ temperature (K). N₂ and O₂ use standard tropospheric mixing ratios; O₃ is initialised at 40 ppb (a typical stratospheric value); O and O(¹D) start at zero."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1000014",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mixing_ratio_to_mol_m3(mixing_ratio, pressure, temperature):\n",
    "    \"\"\"Convert volume mixing ratio to concentration in mol m-3.\"\"\"\n",
    "    return (mixing_ratio * pressure) / (GAS_CONSTANT * temperature)\n",
    "\n",
    "\n",
    "initial_concentrations = {\n",
    "    \"N2\":  mixing_ratio_to_mol_m3(0.78084, pressure, temperature),\n",
    "    \"O2\":  mixing_ratio_to_mol_m3(0.20946, pressure, temperature),\n",
    "    \"O\":   [0.0] * num_grid_cells,\n",
    "    \"O1D\": [0.0] * num_grid_cells,\n",
    "    \"O3\":  mixing_ratio_to_mol_m3(40e-9,   pressure, temperature),  # 40 ppb\n",
    "}\n",
    "\n",
    "state.set_conditions(temperature, pressure)\n",
    "state.set_concentrations(initial_concentrations)\n",
    "\n",
    "# Set photolysis rates\n",
    "user_defined_dict = state.get_user_defined_rate_parameters()\n",
    "for key in user_defined_dict:\n",
    "    user_defined_dict[key] = photolysis_rate_constants[key]\n",
    "state.set_user_defined_rate_parameters(user_defined_dict)\n",
    "\n",
    "print(\"Initial O3 concentrations (mol m-3):\")\n",
    "for alt, conc in zip(vertical_edges, initial_concentrations['O3']):\n",
    "    print(f\"  {alt:.1f} km: {conc:.3e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1000015",
   "metadata": {},
   "source": [
    "## 7. Run the Simulation\n",
    "\n",
    "We integrate forward for 12 hours in 15-minute steps. At the end of each step we:\n",
    "- Advance the simulation clock by the step duration.\n",
    "- Recompute the solar zenith angle for the new time using `pvlib`.\n",
    "- Update the photolysis rates on the state before the next step.\n",
    "\n",
    "This means the photolysis rates change with the diurnal cycle throughout the simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1000016",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_times         = [sim_time]\n",
    "concentrations    = [state.get_concentrations()]\n",
    "photo_rate_history = [{k: v.copy() for k, v in state.get_user_defined_rate_parameters().items()}]\n",
    "\n",
    "time_step         = 15 * 60                    # 15 minutes in seconds\n",
    "simulation_length = 12 * SECONDS_PER_HOUR      # 12 hours in seconds\n",
    "current_time      = 0\n",
    "last_printed_percent = -5\n",
    "\n",
    "while current_time < simulation_length:\n",
    "    # Inner loop: accumulate solver sub-steps until a full time_step is consumed\n",
    "    elapsed = 0\n",
    "    while elapsed < time_step:\n",
    "        remaining = time_step - elapsed\n",
    "        result    = solver.solve(state, remaining)\n",
    "        elapsed      += result.stats.final_time\n",
    "        current_time += result.stats.final_time\n",
    "        if result.state != SolverState.Converged:\n",
    "            print(f\"  Solver state: {result.state} at t={current_time:.0f} s\")\n",
    "\n",
    "    # Advance the wall-clock time and update photolysis rates for the new position\n",
    "    sim_time += timedelta(seconds=time_step)\n",
    "    photolysis_rate_constants, _ = get_tuv_rates(sim_time, num_grid_cells, start)\n",
    "\n",
    "    user_defined_dict = state.get_user_defined_rate_parameters()\n",
    "    for key in user_defined_dict:\n",
    "        user_defined_dict[key] = photolysis_rate_constants[key]\n",
    "    state.set_user_defined_rate_parameters(user_defined_dict)\n",
    "\n",
    "    sim_times.append(sim_time)\n",
    "    concentrations.append(state.get_concentrations())\n",
    "    photo_rate_history.append({k: v.copy() for k, v in state.get_user_defined_rate_parameters().items()})\n",
    "\n",
    "    current_percent = (current_time / simulation_length) * 100\n",
    "    rounded = int(current_percent // 5) * 5\n",
    "    if rounded > last_printed_percent:\n",
    "        last_printed_percent = rounded\n",
    "        local_str = sim_time.astimezone(boulder_tz).strftime('%H:%M %Z')\n",
    "        print(f\"Progress: {last_printed_percent}%  (sim time: {local_str})\")\n",
    "\n",
    "print(\"Simulation complete.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1000017",
   "metadata": {},
   "source": [
    "## 8. Build the xarray Dataset\n",
    "\n",
    "We assemble results into an `xr.Dataset`. Unlike the TS1 tutorial, time is stored as actual `datetime64` timestamps (not elapsed seconds) and photolysis rates are recorded at every time step so we can see how they vary with the diurnal cycle."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1000018",
   "metadata": {},
   "outputs": [],
   "source": [
    "final_conditions = state.get_conditions()\n",
    "\n",
    "data_vars = {\n",
    "    \"temperature\": ([\"height\"], final_conditions[\"temperature\"], {\"units\": \"K\"}),\n",
    "    \"pressure\":    ([\"height\"], final_conditions[\"pressure\"],    {\"units\": \"Pa\"}),\n",
    "    \"air_density\": ([\"height\"], final_conditions[\"air_density\"], {\"units\": \"mol m-3\"}),\n",
    "}\n",
    "\n",
    "# Photolysis rates over time\n",
    "photo_names  = list(photo_rate_history[0].keys())\n",
    "photo_values = [[params[name] for name in photo_names] for params in photo_rate_history]\n",
    "data_vars[\"photolysis_rates\"] = (\n",
    "    [\"time\", \"photolysis_rate_names\", \"height\"],\n",
    "    photo_values,\n",
    "    {\"units\": \"s-1\"}\n",
    ")\n",
    "\n",
    "# Species concentrations\n",
    "for species in concentrations[0]:\n",
    "    data_vars[species] = (\n",
    "        [\"time\", \"height\"],\n",
    "        np.array([c[species] for c in concentrations], dtype=np.float64),\n",
    "        {\"units\": \"mol m-3\"}\n",
    "    )\n",
    "\n",
    "ds = xr.Dataset(\n",
    "    data_vars,\n",
    "    coords={\n",
    "        \"time\":                  np.array(sim_times, dtype=\"datetime64[s]\"),\n",
    "        \"height\":                vertical_edges,\n",
    "        \"photolysis_rate_names\": np.array(photo_names, dtype=\"S\"),\n",
    "    }\n",
    ")\n",
    "\n",
    "print(ds)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1000019",
   "metadata": {},
   "source": [
    "## 9. Plot Results\n",
    "\n",
    "We plot the evolution of O, O(¹D), and O₃ over the 12-hour simulation. Each line is a different altitude grid cell. The diurnal signature in the photolysis rates should be visible in the O and O(¹D) concentrations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1000020",
   "metadata": {},
   "outputs": [],
   "source": [
    "elapsed_hours = (ds['time'] - ds['time'].isel(time=0)) / np.timedelta64(1, 'h')\n",
    "\n",
    "fig = plt.figure(figsize=(12, 8))\n",
    "gs  = fig.add_gridspec(2, 2)\n",
    "\n",
    "ax_o   = fig.add_subplot(gs[0, 0])\n",
    "ax_o1d = fig.add_subplot(gs[0, 1])\n",
    "ax_o3  = fig.add_subplot(gs[1, :])\n",
    "\n",
    "for cell in range(num_grid_cells):\n",
    "    label = f\"{vertical_edges[cell]:.1f} km\"\n",
    "    ax_o.plot(elapsed_hours,   ds['O'].isel(height=cell),   label=label)\n",
    "    ax_o1d.plot(elapsed_hours, ds['O1D'].isel(height=cell), label=label)\n",
    "    ax_o3.plot(elapsed_hours,  ds['O3'].isel(height=cell),  label=label)\n",
    "\n",
    "for ax, title in [(ax_o, 'O'), (ax_o1d, 'O1D'), (ax_o3, 'O3')]:\n",
    "    ax.set_title(title)\n",
    "    ax.set_xlabel('Time [hours since start]')\n",
    "    ax.set_ylabel('Concentration [mol m⁻³]')\n",
    "    ax.set_xlim(0, simulation_length / SECONDS_PER_HOUR)\n",
    "    ax.set_ylim(0, None)\n",
    "    ax.grid(True, alpha=0.5)\n",
    "    ax.spines[:].set_visible(False)\n",
    "    ax.tick_params(width=0)\n",
    "\n",
    "handles, labels = ax_o3.get_legend_handles_labels()\n",
    "ax_o3.legend(handles, labels, title='Altitude', ncol=5, loc='upper right')\n",
    "\n",
    "fig.tight_layout()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv (3.12.3)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}